package com.ruoyi.mqtt.aspect;

import com.ruoyi.mqtt.annotation.QpsMonitor;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;

@Aspect
@Component
@Slf4j
public class QpsMonitorAspect {
    private final Map<String, QpsCounter> counterMap = new ConcurrentHashMap<>();

    @Around("@annotation(qpsMonitor)")
    public Object around(ProceedingJoinPoint point, QpsMonitor qpsMonitor) throws Throwable {
        String monitorName = qpsMonitor.value().isEmpty() ? 
            point.getSignature().toShortString() : qpsMonitor.value();
        
        QpsCounter counter = counterMap.computeIfAbsent(monitorName,
            k -> new QpsCounter(qpsMonitor.timeWindow()));

        try {
            counter.increment();
            return point.proceed();
        } finally {
            if (counter.shouldPrintStats()) {
                double qps = counter.calculateQps();
                log.info("[QPS统计] {}: {}/s", monitorName, String.format("%.2f", qps));
                counter.resetIfNeeded();
            }
        }
    }

    private static class QpsCounter {
        private final AtomicLong count = new AtomicLong(0);
        private final AtomicLong lastPrintTime = new AtomicLong(System.currentTimeMillis());
        private final int timeWindowSeconds;

        public QpsCounter(int timeWindowSeconds) {
            this.timeWindowSeconds = timeWindowSeconds;
        }

        public void increment() {
            count.incrementAndGet();
        }

        public boolean shouldPrintStats() {
            return System.currentTimeMillis() - lastPrintTime.get() >= timeWindowSeconds * 1000;
        }

        public double calculateQps() {
            long currentTime = System.currentTimeMillis();
            long timeElapsed = currentTime - lastPrintTime.get();
            return count.get() * 1000.0 / timeElapsed;
        }

        public void resetIfNeeded() {
            count.set(0);
            lastPrintTime.set(System.currentTimeMillis());
        }
    }
} 